import docx2txt
import PyPDF2
import requests
from bs4 import BeautifulSoup
import os

class DocumentService:
    def __init__(self):
        self.db_path = os.path.join("data", "documents.db")
        os.makedirs("data", exist_ok=True)

    def load_document(self, path):
        """加载文档内容"""
        try:
            file_ext = os.path.splitext(path)[1].lower()
            
            if file_ext == '.txt':
                with open(path, 'r', encoding='utf-8') as f:
                    return f.read()
                    
            elif file_ext == '.docx':
                return docx2txt.process(path)
                
            elif file_ext == '.pdf':
                text = ""
                with open(path, 'rb') as f:
                    pdf_reader = PyPDF2.PdfReader(f)
                    for page in pdf_reader.pages:
                        text += page.extract_text()
                return text
                
            else:
                raise ValueError(f"不支持的文件类型: {file_ext}")
                
        except UnicodeDecodeError:
            # 如果 UTF-8 解码失败，尝试其他编码
            try:
                with open(path, 'r', encoding='gbk') as f:
                    return f.read()
            except:
                try:
                    with open(path, 'r', encoding='gb2312') as f:
                        return f.read()
                except:
                    raise Exception(f"无法识别文件编码: {path}")
        except Exception as e:
            raise Exception(f"处理文档时出错: {str(e)}")
    
    def load_url(self, url):
        """从网页加载内容"""
        try:
            response = requests.get(url)
            response.raise_for_status()
            
            soup = BeautifulSoup(response.text, 'html.parser')
            
            # 移除脚本和样式
            for script in soup(["script", "style"]):
                script.decompose()
                
            # 获取文本
            text = soup.get_text()
            
            # 清理文本
            lines = (line.strip() for line in text.splitlines())
            chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
            text = ' '.join(chunk for chunk in chunks if chunk)
            
            return text
            
        except Exception as e:
            raise Exception(f"加载网页失败: {str(e)}") 